`timescale 1ns / 1ps

// kyber compress
//
// kyber512/768 Compress0, du = 10, dv = 4
// Compress0 -- du = 10, transform 12bit to 10bit
// 			 -- dv =  4, transform 12bit to  4bit
//
// kyber1024 Compress1, du = 11, dv = 5
// Compress1 -- du = 11, transform 12bit to 11bit
// 			 -- dv =  5, transform 12bit to  5bit

`define	k_width 12

`define	du10	10
`define	dv4		4

`define	du11	11
`define	dv5		5

`define coeff_num 1 // coefficient number


module Kyber_Compress(

	input 	d_mode, // d_mode, 0 = du, 1 = dv
	input 	k_mode, // k_mode, 0 = kyber512/768, 1 = kyber1024
					// for d_mode and k_mode select, the state is below 
					//		  0	        0            du10,	du=10 kyber512/768
					//		  1	        0            dv4,	dv=4  kyber512/768
					//		  0	        1            du11,	du=11 kyber1024
					//		  1	        1            dv5,	dv=5  kyber1024

	output 	[22*`coeff_num-1:0]	Mul_A,
	output 	[12*`coeff_num-1:0]	Mul_B,
	input 	[34*`coeff_num-1:0]	Mul_P,

	input 	[`k_width*`coeff_num - 1 : 0] data_in,
	output 	[`k_width*`coeff_num - 1 : 0] data_out

    );

parameter	index_du = 22'd 2580335 ; 	// for du=10 / du=11, they are the same index
parameter	index_dv = 22'd 315 ;		// for dv=4 / dv=5, they are the same index

wire	[`du10*`coeff_num - 1 : 0]	do_du10;
wire	[`dv4*`coeff_num  - 1 : 0]	do_dv4;
wire	[`du11*`coeff_num - 1 : 0]	do_du11;
wire	[`dv5*`coeff_num  - 1 : 0]	do_dv5;

wire	[(`du10+3)*`coeff_num - 1 : 0]	do_du10_temp;
wire	[(`dv4+3)*`coeff_num  - 1 : 0]	do_dv4_temp;
wire	[(`du11+3)*`coeff_num - 1 : 0]	do_du11_temp;
wire	[(`dv5+3)*`coeff_num  - 1 : 0]	do_dv5_temp;

reg 	[`k_width*`coeff_num - 1 : 0] data_out_temp;

assign 	data_out = data_out_temp;

genvar i;

generate
	for (i = 0; i < `coeff_num; i = i + 1)
	begin
		assign Mul_A[22*i+21:22*i] = d_mode ? index_dv : index_du;

		assign do_du11_temp[(`du11+3)*i+(`du11+3)-1:(`du11+3)*i] = (Mul_P[34*i+33:34*i+21] + 1'd1);
		assign do_du11[`du11*i+`du11-1:`du11*i] = do_du11_temp[(`du11+3)*i+(`du11+3)-1:(`du11+3)*i+1];

		assign do_dv5_temp[(`dv5+3)*i+(`dv5+3)-1:(`dv5+3)*i] = (Mul_P[34*i+33:34*i+14] + 1'd1);
		assign do_dv5[`dv5*i+`dv5-1:`dv5*i] = do_dv5_temp[(`dv5+3)*i+(`dv5+3)-1:(`dv5+3)*i+1];

		assign do_du10_temp[(`du10+3)*i+(`du10+3)-1:(`du10+3)*i] = (Mul_P[34*i+33:34*i+22] + 1'd1);
		assign do_du10[`du10*i+`du10-1:`du10*i] = do_du10_temp[(`du10+3)*i+(`du10+3)-1:(`du10+3)*i+1];

		assign do_dv4_temp[(`dv4+3)*i+(`dv4+3)-1:(`dv4+3)*i] = (Mul_P[34*i+33:34*i+15] + 1'd1);
		assign do_dv4[`dv4*i+`dv4-1:`dv4*i] = do_dv4_temp[(`dv4+3)*i+(`dv4+3)-1:(`dv4+3)*i+1];

		// assign do_du_temp[(`du+2)*i+(`du+2)-1:(`du+2)*i] = (Mul_P_T2[34*i+33:34*i+22] + 1'd1);
		// assign do_du[`du*i+`du-1:`du*i] = do_du_temp[(`du+2)*i+(`du+2)-1:(`du+2)*i+1] ;

		// assign do_dv_temp[(`du-3)*i+(`du-3)-1:(`du-3)*i] = (Mul_P_T2[34*i+20:34*i+15] + 1'd1);
		// //assign do_dv[`du*i+`du-1:`du*i] = { 6'd0,do_dv_temp[(`du-3)*i+(`du-3)-1:(`du-3)*i+1] };
		// assign do_dv[`du*i+`du-1:`du*i] = { 6'd0,do_dv_temp[(`du-3)*i+4:(`du-3)*i+1] };
	end
endgenerate

generate
	for ( i = 0; i < `coeff_num; i = i + 1 ) begin
		always@(*)
		begin 
			case ({d_mode,k_mode})
				2'b00:	data_out_temp[`k_width*i + `k_width - 1:`k_width*i] = {2'd0,do_du10[`du10*i + `du10 - 1:`du10*i]};	// du-10
				2'b10:	data_out_temp[`k_width*i + `k_width - 1:`k_width*i] = {8'd0,do_dv4[`dv4*i + `dv4 - 1:`dv4*i]};		// dv=4
				2'b01:  data_out_temp[`k_width*i + `k_width - 1:`k_width*i] = {1'd0,do_du11[`du11*i + `du11 - 1:`du11*i]};	// du=11
				2'b11:  data_out_temp[`k_width*i + `k_width - 1:`k_width*i] = {7'd0,do_dv5[`dv5*i + `dv5 - 1:`dv5*i]};		// dv=5
			endcase
		end
	end
endgenerate


assign Mul_B = data_in;


endmodule
